import json
import random
import re
from typing import List, Dict, Any, Tuple
import copy


class MeetingScheduleEnv:
    """智能会议安排环境"""

    def __init__(self, max_steps=10, **kwargs):
        """
        初始化会议安排环境
        
        Args:
            max_steps: 最大交互步数
        """
        self.max_steps = max_steps
        self._step = 0
        
        # 当前任务状态
        self.meeting_topic = None
        self.required_attendees = []
        self.optional_attendees = []
        self.time_range = []
        self.availability = {}
        self.historical_pattern = {}
        self.negotiation_history = []
        self.current_selected_time = None
        
        # 模拟的MCP工具调用结果缓存
        self.calendar_cache = {}
        self.history_cache = {}
        
        # 任务完成标志
        self.task_done = False
        
    def reset(self, task_data: Dict) -> Tuple[str, Dict]:
        """
        重置环境，开始新的会议安排任务
        
        Args:
            task_data: 包含会议信息的字典
                - meeting_topic: 会议主题
                - required_attendees: 必选参会人员列表
                - optional_attendees: 可选参会人员列表
                - availability: 人员日程可用性
                - historical_pattern: 历史会议模式
                
        Returns:
            observation: 初始观察（任务描述）
            info: 额外信息
        """
        self._step = 0
        self.task_done = False
        self.negotiation_history = []
        self.current_selected_time = None
        
        # 解析任务数据
        self.meeting_topic = task_data['meeting_topic']
        self.required_attendees = task_data['required_attendees']
        self.optional_attendees = task_data.get('optional_attendees', [])
        self.availability = task_data['availability']
        self.historical_pattern = task_data.get('historical_pattern', {})
        
        # 生成时间范围
        days = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
        time_slots = ['morning', 'afternoon', 'evening']
        self.time_range = [(day, slot) for day in days for slot in time_slots]
        
        # 构建初始观察
        observation = self._build_initial_observation()
        
        info = {
            'required_attendees': self.required_attendees,
            'optional_attendees': self.optional_attendees,
            'meeting_topic': self.meeting_topic
        }
        
        return observation, info
    
    def _build_initial_observation(self) -> str:
        """构建初始观察文本"""
        obs = f"会议安排任务：{self.meeting_topic}\n"
        obs += f"必选参会人员：{', '.join(self.required_attendees)}\n"
        obs += f"可选参会人员：{', '.join(self.optional_attendees)}\n"
        obs += f"时间范围：下周（周一至周日，分上午/下午/晚上三个时段）\n"
        obs += f"请使用MCP工具来查询日历、历史会议模式，并安排最优的会议时间。"
        return obs
    
    def step(self, action: Dict, action_valid: bool) -> Tuple[str, float, bool, Dict]:
        """
        执行一个动作
        
        Args:
            action: 动作字典，包含tool_name和tool_input
            action_valid: 动作是否有效
            
        Returns:
            observation: 执行结果的观察
            reward: 即时奖励
            done: 是否完成任务
            info: 额外信息
        """
        self._step += 1
        
        if not action_valid:
            return "动作格式无效，请检查格式。", -10, True, {}
        
        # 达到最大步数
        if self._step >= self.max_steps:
            if not self.task_done:
                return "已达到最大步数，任务失败。", -50, True, {}
        
        # 执行动作
        tool_name = action.get('tool_name', '')
        tool_input = action.get('tool_input', {})
        
        observation, reward, done = self._execute_tool(tool_name, tool_input)
        
        info = {
            'step': self._step,
            'tool_used': tool_name,
            'negotiations_count': len(self.negotiation_history)
        }
        
        return observation, reward, done, info
    
    def _execute_tool(self, tool_name: str, tool_input: Dict) -> Tuple[str, float, bool]:
        """
        执行MCP工具调用
        
        Args:
            tool_name: 工具名称
            tool_input: 工具输入参数
            
        Returns:
            observation: 工具执行结果
            reward: 即时奖励
            done: 是否完成任务
        """
        reward = 0
        done = False
        
        if tool_name == 'query_calendar':
            # 查询日历
            observation = self._query_calendar(tool_input)
            reward = 0.5  # 查询动作有小额正奖励
            
        elif tool_name == 'query_history':
            # 查询历史会议模式
            observation = self._query_history(tool_input)
            reward = 0.5
            
        elif tool_name == 'negotiate_time':
            # IM协商
            observation = self._negotiate_time(tool_input)
            reward = -2  # 协商有成本
            
        elif tool_name == 'select_time':
            # 选择会议时间（最终决策）
            observation, final_reward, done = self._select_time(tool_input)
            reward = final_reward
            self.task_done = done
            
        elif tool_name == 'calculate':
            # 计算器
            observation = self._calculate(tool_input)
            reward = 0
            
        else:
            observation = f"未知的工具: {tool_name}"
            reward = -5
        
        return observation, reward, done
    
    def _query_calendar(self, tool_input: Dict) -> str:
        """查询指定人员的日历"""
        persons = tool_input.get('persons', [])
        time_range = tool_input.get('time_range', 'next_week')
        
        if not persons:
            return "错误：未指定查询人员"
        
        result = f"日历查询结果（{time_range}）：\n"
        for person in persons:
            if person not in self.availability:
                result += f"{person}: 无数据\n"
                continue
            
            result += f"\n{person}:\n"
            person_schedule = self.availability[person]
            
            for day, slots in person_schedule.items():
                result += f"  {day}:\n"
                for slot, status in slots.items():
                    result += f"    {slot}: {status}\n"
        
        return result
    
    def _query_history(self, tool_input: Dict) -> str:
        """查询历史会议模式"""
        meeting_topic = tool_input.get('meeting_topic', self.meeting_topic)
        
        if not self.historical_pattern:
            return f"{meeting_topic} 暂无历史数据"
        
        pattern = self.historical_pattern
        result = f"{meeting_topic} 历史会议模式：\n"
        result += f"最常见时间：{pattern.get('most_common_day', '未知')}, {pattern.get('most_common_slot', '未知')}\n"
        result += f"频率：{pattern.get('frequency', 0)}%\n"
        
        if 'recent_meetings' in pattern:
            result += "\n最近5次会议：\n"
            for meeting in pattern['recent_meetings']:
                result += f"  Week {meeting['week']}: {meeting['day']}, {meeting['slot']}\n"
        
        return result
    
    def _negotiate_time(self, tool_input: Dict) -> str:
        """与人员协商时间"""
        person = tool_input.get('person', '')
        day = tool_input.get('day', '')
        time_slot = tool_input.get('time_slot', '')
        
        if not all([person, day, time_slot]):
            return "错误：协商参数不完整"
        
        # 模拟协商成功（简化处理）
        success = True
        
        if success:
            # 更新可用性
            if person in self.availability and day in self.availability[person]:
                self.availability[person][day][time_slot] = 'available'
            
            # 记录协商历史
            self.negotiation_history.append({
                'person': person,
                'day': day,
                'time_slot': time_slot,
                'result': 'success'
            })
            
            return f"协商成功：{person} 已同意在 {day} {time_slot} 参加会议"
        else:
            return f"协商失败：{person} 无法在 {day} {time_slot} 参加"
    
    def _select_time(self, tool_input: Dict) -> Tuple[str, float, bool]:
        """选择最终会议时间"""
        day = tool_input.get('day', '')
        time_slot = tool_input.get('time_slot', '')
        
        if not day or not time_slot:
            return "错误：时间选择参数不完整", -20, False
        
        self.current_selected_time = {'day': day, 'time_slot': time_slot}
        
        # 计算最终奖励
        final_reward = self._calculate_final_reward(day, time_slot)
        
        observation = f"已选择会议时间：{day} {time_slot}\n"
        observation += f"最终奖励评分：{final_reward:.2f}"
        
        return observation, final_reward, True
    
    def _calculate_final_reward(self, day: str, time_slot: str) -> float:
        """计算最终奖励分数"""
        reward = 0
        
        # 统计参与情况
        required_present = 0
        optional_present = 0
        
        # 检查必选人员
        for person in self.required_attendees:
            if person in self.availability:
                if day in self.availability[person]:
                    if self.availability[person][day].get(time_slot, 'busy') == 'available':
                        required_present += 1
        
        # 检查可选人员
        for person in self.optional_attendees:
            if person in self.availability:
                if day in self.availability[person]:
                    if self.availability[person][day].get(time_slot, 'busy') == 'available':
                        optional_present += 1
        
        # 1. 必选人员参与率（权重最高）
        required_rate = required_present / max(1, len(self.required_attendees))
        if required_rate < 1.0:
            missing = len(self.required_attendees) - required_present
            reward -= missing * 100
        
        # 2. 可选人员参与率
        optional_rate = optional_present / max(1, len(self.optional_attendees))
        reward += optional_present * 1
        
        # 3. 协商成本
        negotiations_cnt = len(self.negotiation_history)
        reward -= negotiations_cnt * 2
        
        # 4. 历史模式匹配
        if self.historical_pattern:
            most_common_day = self.historical_pattern.get('most_common_day', '')
            most_common_slot = self.historical_pattern.get('most_common_slot', '')
            
            if day == most_common_day and time_slot == most_common_slot:
                frequency = self.historical_pattern.get('frequency', 0) / 100.0
                reward += frequency * 5
        
        return reward
    
    def _calculate(self, tool_input: Dict) -> str:
        """计算器工具"""
        expression = tool_input.get('expression', '')
        
        try:
            # 安全的计算（仅支持基本运算）
            result = eval(expression, {"__builtins__": {}}, {})
            return f"计算结果：{expression} = {result}"
        except Exception as e:
            return f"计算错误：{str(e)}"
    
    @property
    def action_space(self):
        return NotImplementedError
    
    @property
    def observation_space(self):
        return NotImplementedError
    
    def close(self):
        """关闭环境"""
        pass

